import pandas as pd
import numpy as np
import pickle
import math


class PreHandle:

    def __init__(self):
        pass

    def prehandle(self, dict):
        # 对原始数据的副本去除前7年左右所有的那三只股票，如果那三只股票存在缺失数据，就把那三只股票全都去掉
        key_list = sorted(dict.keys())
        for data in key_list:
            series = dict[data]['open']
            if math.isnan(series['002558.XSHE']) or math.isnan(series['603110.XSHG']) or math.isnan(series['002611.XSHE']):
                dfT = dict[data].T
                del dfT['002558.XSHE']
                del dfT['603110.XSHG']
                del dfT['002611.XSHE']
                dict[data] = dfT.T

    def prehandle_db_avg_stgy(self,dict):
        # 对原始数据的副本去除近9年来的那三只股票的数据
        key_list = sorted(dict.keys())
        for data in key_list:
            dfT = dict[data].T
            del dfT['002558.XSHE']
            del dfT['603110.XSHG']
            del dfT['002611.XSHE']
            dict[data] = dfT.T
                